package mr

import (
	"log"
	"strconv"
	"sync"
	"time"
)
import "net"
import "os"
import "net/rpc"
import "net/http"

const ReduceOutput = "mr-out-"
const ReduceInput = "mr-tmp-"

type Task struct {
	workerId int
	stage    bool
	deadline time.Time
}

type Coordinator struct {
	// 0代表任务完成,1代表map阶段,2代表reduce阶段
	stage           int
	toMap           chan string
	toReduce        chan string
	nReduce         int
	mapWaitGroup    sync.WaitGroup
	reduceWaitGroup sync.WaitGroup
	mapTasks        map[string]*Task
	reduceTasks     map[string]*Task
	// 用于保护整个结构
	lock sync.Mutex
}

// 与worker通信

func (c *Coordinator) GetTask(args *GetTaskArgs, reply *GetTaskReply) (err error) {
	c.lock.Lock()
	reply.Stage = c.stage
	reply.NReduce = c.nReduce
	if c.stage == 1 {
		if len(c.toMap) == 0 {
			reply.Stage = 4
		} else {
			task := <-c.toMap
			reply.InputFile = task
			c.mapTasks[task].workerId = args.WorkerId
			c.mapTasks[task].deadline = time.Now().Add(time.Second * 10)
		}
	} else if c.stage == 2 {
		if len(c.toReduce) == 0 {
			reply.Stage = 4
		} else {
			task := <-c.toReduce
			reply.InputFile = task
			c.reduceTasks[task].workerId = args.WorkerId
			c.reduceTasks[task].deadline = time.Now().Add(time.Second * 10)
		}
	}
	c.lock.Unlock()
	return
}

func (c *Coordinator) FinishTask(args *FinishTaskArgs, reply *FinishTaskReply) (err error) {
	c.lock.Lock()
	if c.stage == 1 {
		if c.mapTasks[args.InputFile].workerId == args.WorkerId {
			c.mapTasks[args.InputFile].stage = true
			reply.Ok = true
			//fmt.Printf("%v 任务完成，允许写入文件\n", args.InputFile)
		} else {
			reply.Ok = false
		}
	} else if c.stage == 2 {
		if c.reduceTasks[args.InputFile].workerId == args.WorkerId {
			c.reduceTasks[args.InputFile].stage = true
			reply.Ok = true
			//fmt.Printf("%v 任务完成，允许写入文件\n", args.InputFile)
		} else {
			reply.Ok = false
		}
	}
	c.lock.Unlock()
	return
}

// 提交任务
func (c *Coordinator) CommitTask(args *CommitTaskArgs, reply *CommitTaskReply) (err error) {
	c.lock.Lock()
	st := c.stage
	c.lock.Unlock()
	if st == 1 {
		c.mapWaitGroup.Done()
	} else {
		c.reduceWaitGroup.Done()
	}
	reply.Ok = true
	return
}

// 超时则重新分配任务

func (c *Coordinator) CheckTimeOut() {
	for !c.Done() {
		time.Sleep(time.Second)
		c.lock.Lock()
		if c.stage == 1 {
			for file, task := range c.mapTasks {
				if task.workerId == 0 || task.stage || time.Now().Before(task.deadline) {
					continue
				} else {
					c.toMap <- file
					task.workerId = 0
				}
			}
		} else {
			for file, task := range c.reduceTasks {
				if task.workerId == 0 || task.stage || time.Now().Before(task.deadline) {
					continue
				} else {
					c.toReduce <- file
					task.workerId = 0
				}
			}
		}
		c.lock.Unlock()
	}
}

// 查看任务是否并改变stage

func (c *Coordinator) ChangeStage() {
	c.mapWaitGroup.Wait()
	c.lock.Lock()
	c.stage = 2
	c.lock.Unlock()
	c.reduceWaitGroup.Wait()
	c.lock.Lock()
	c.stage = 0
	c.lock.Unlock()
}

// start a thread that listens for RPCs from worker.go
func (c *Coordinator) server() {
	rpc.Register(c)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := coordinatorSock()
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

// Done
// main/mrcoordinator.go calls Done() periodically to find out
// if the entire job has finished.
func (c *Coordinator) Done() bool {
	c.lock.Lock()
	t := c.stage
	c.lock.Unlock()
	if t == 0 {
		return true
	} else {
		return false
	}
}

// MakeCoordinator
// create a Coordinator.
// main/mrcoordinator.go calls this function.
// nReduce is the number of reduce tasks to use.
func MakeCoordinator(files []string, nReduce int) *Coordinator {
	c := Coordinator{
		stage:       1,
		nReduce:     nReduce,
		toMap:       make(chan string, len(files)),
		toReduce:    make(chan string, nReduce),
		mapTasks:    make(map[string]*Task),
		reduceTasks: make(map[string]*Task),
	}
	c.mapWaitGroup.Add(len(files))
	c.reduceWaitGroup.Add(nReduce)
	//fmt.Println("需要处理的文件有")
	//fmt.Println("------------")
	for _, file := range files {
		//fmt.Println(file)
		c.toMap <- file
		c.mapTasks[file] = new(Task)
	}
	//fmt.Println("------------")
	for i := 0; i < nReduce; i++ {
		reduceInputFile := strconv.Itoa(i)
		c.toReduce <- reduceInputFile
		c.reduceTasks[reduceInputFile] = new(Task)
		os.Create(ReduceInput + strconv.Itoa(i))
	}
	go c.ChangeStage()
	go c.CheckTimeOut()
	c.server()
	return &c
}
